"""
Author: Morphlng
Date: 2024-03-18 16:50:51
LastEditTime: 2025-02-23 10:02:58
LastEditors: Morphlng
Description: Test for creating scenario in Carla
FilePath: /DrivingGym/tests/carla/carla_scenario_test.py
"""

from driving_gym.environment.scenario import ScenarioManager
from driving_gym.environment.scenario.config import CarlaScenarioConfig, RoutingTask
from driving_gym.simulation import AdapterFactory
from driving_gym.simulation.adapter.carla import (
    CarlaActorConfig,
    CarlaConfig,
    CarlaSensorConfig,
)

carla_config = (
    CarlaConfig()
    .connection(host="localhost", port=2000)
    .synchronization(sync_mode=False, fixed_delta_seconds=None)
)

sensor_config = (
    CarlaSensorConfig()
    .blueprint(
        type="bev",
        attributes={
            "target_size": (512, 512),
            "render_lanes_on_junctions": True,
        },
    )
    .spawning(parent="hero")
)

hero_config = (
    CarlaActorConfig()
    .blueprint("vehicle.tesla.model3", rolename="hero")
    .spawning((252.4, 195.2, 0.5))
    .add_sensor("bev", sensor_config.to_dict())
)

hero_task = RoutingTask().starting((252.4, 195.2, 0.5)).ending((187.52, 195.20, 0.5))

scenario = (
    CarlaScenarioConfig()
    .environment(map_name="Town01")
    .add_actor("hero", hero_config.to_dict())
    .add_task("hero", hero_task.to_dict())
    .traffic(npc_vehicle_num=10, npc_walker_num=5, spawn_near_agents=True)
)

adapter = AdapterFactory.create(carla_config.to_dict())

scenario_manager = ScenarioManager(scenario.to_dict(), adapter)


def test_adapter_initialize():
    assert "CarlaAdapter" in str(type(adapter)), "Adapter type should be CarlaAdapter"
    adapter.initialize()
    assert adapter._client is not None, "Client should not be None"


def test_load_scenario():
    scenario = scenario_manager.get_scenario()
    scenario.reset()

    hero_location = adapter.get_transform("hero").location
    spawn_point = CarlaActorConfig.parse_location(
        hero_config.actor_spawn_point
    ).location
    assert hero_location.distance(spawn_point) < 1, "Hero should be at the spawn point"

    assert (
        len(adapter._traffic_vehicles) + len(adapter._traffic_walkers) > 0
    ), "Traffic should be generated"


def test_bev_data():
    bev_data = adapter.get_sensor_data("hero:bev")
    assert bev_data is not None, "BEV data should not be None"

    img_x = sensor_config.sensor_attributes["target_size"][0]
    img_y = sensor_config.sensor_attributes["target_size"][1]

    assert bev_data.shape == (
        img_x,
        img_y,
        3,
    ), f"BEV data shape should be ({img_x}, {img_y}, 3)"


def test_load_another_scenario():
    traffic_vehicles = adapter._traffic_vehicles.copy()
    traffic_walkers = adapter._traffic_walkers.copy()

    scenario = scenario_manager.get_scenario()
    scenario.reset()

    # Traffic should be destroyed then regenerated
    for vehicle in traffic_vehicles.values():
        assert (
            not vehicle.is_alive or not vehicle.is_active
        ), "Traffic vehicles should be destroyed"
    for walker in traffic_walkers.values():
        assert (
            not walker.is_alive or not walker.is_active
        ), "Traffic walkers should be destroyed"

    # New sceanrio should be loaded
    hero_location = adapter.get_transform("hero").location
    spawn_point = CarlaActorConfig.parse_location(
        hero_config.actor_spawn_point
    ).location
    assert hero_location.distance(spawn_point) < 1, "Hero should be at the spawn point"

    assert (
        len(adapter._traffic_vehicles) + len(adapter._traffic_walkers) > 0
    ), "Traffic should be generated"


def test_cleanup():
    adapter.cleanup()
    assert adapter._client is None, "Client should be None"
